package assignment04;

import gnu.trove.iterator.TIntIterator;
import gnu.trove.list.TIntList;

import java.util.List;
import java.util.Map;

import nlp.util.Indexer;

import assignment04.WordAlignmentTester.Alignment;
import assignment04.WordAlignmentTester.SentencePair;
import assignment04.WordAlignmentTester.WordAligner;

public class HMMWordAligner implements WordAligner {
	private final static int NULL_INDEX = -1;
	private final static double PNULL = .2;
	private final static double PREST = .8;
	
	public static class HMMFactory implements WordAlignerFactory {
		IntCounter frenchCounter;
		IntCounter englishCounter;
		IntCounterMap counterMap;
		IntCounterMap tMap;
		Indexer<String> englishIndexer;
		Indexer<String> frenchIndexer;
		
		IntCounter jumpCounter;
		Indexer<TIntList> positionIndexer;
		
		List<SentencePair> validatePairs1;
		Map<Integer,Alignment> validateAlignments_1;
		List<SentencePair> validatePairs2;
		Map<Integer,Alignment> validateAlignments_2;
		
		public WordAligner trainAligner(List<SentencePair> trainingSentencePairs) {
			buildIndexerAndCounter(trainingSentencePairs);
			// build jump width model
			buildJumpCounter();
			System.out.println("jumpCounter size is " + jumpCounter.size());

			for (int i = 0; i < 5; i++)
			{
				System.err.println("model one training loop " + i);
				modelOneTrainEStep(trainingSentencePairs);
				modelOneTrainMStep();
			}
			return new HMMWordAligner(frenchCounter, englishCounter, tMap, englishIndexer, frenchIndexer, jumpCounter, positionIndexer);
			
		}
		
		

		private void modelOneTrainMStep() {
			tMap = IntCounters.conditionalNormalize(counterMap);
			counterMap.reset();
		}
		
		private void buildIndexerAndCounter(List<SentencePair> trainingSentencePairs) {
			for (SentencePair pair : trainingSentencePairs) {
				List<String> eList = pair.englishWords;
				List<String> fList = pair.frenchWords;

				for (String eWord : eList) {
					englishCounter.incrementCount(englishIndexer.addAndGetIndex(eWord), 1.0);
				}
				
				for (String fWord : fList) {
					int fIndex = frenchIndexer.addAndGetIndex(fWord);
					frenchCounter.incrementCount(fIndex, 1.0);
					counterMap.setCount(NULL_INDEX, fIndex, 0.0);
				}
				
				for (String english : eList) {
					int eIndex = englishIndexer.indexOf(english);
					for (String french : fList) {
						int fIndex = frenchIndexer.indexOf(french);
						counterMap.setCount(eIndex, fIndex, 0.0);
					}
				}
			}
			
			
			//normalize
			englishCounter.normalize();
			frenchCounter.normalize();
			
			
			System.out.println("there are " + englishCounter.size() + " english word types");
			System.out.println("there are " + frenchCounter.size() + " french word types");
			
			for (TIntIterator outerIt = counterMap.keySet().iterator(); outerIt.hasNext();) {
				int key = outerIt.next();
				IntCounter counter = counterMap.getCounter(key);
				for (TIntIterator innerIt = counter.keySet().iterator(); innerIt.hasNext();) {
					int value = innerIt.next();
					tMap.setCount(key, value, frenchCounter.getCount(value));
				}
			}
		}
		
		private void buildJumpCounter() {
			for (SentencePair pair : validatePairs1) {
				innerBuildJumpCounter(pair, validateAlignments_1);
			}
			
			for (SentencePair pair : validatePairs2) {
				innerBuildJumpCounter(pair, validateAlignments_2);
			} 
		}
		
		private void innerBuildJumpCounter(SentencePair pair, Map<Integer, Alignment> alignments) {
			 Alignment referenceAlignment = alignments.get(pair.getSentenceID());
		      if (referenceAlignment == null) {
		    	  return;
		      }

			 int englishLength = pair.getEnglishWords().size();
//			 System.out.println("pair " + pair.getSentenceID() + " french size : " + pair.frenchWords.size());
			 for (int frenchPosition = 1, previous = 0, current = 0; frenchPosition <= pair.getFrenchWords().size(); frenchPosition++) {
				 for (int englishPosition = 0; englishPosition <= englishLength; englishPosition++) {
			        	if (referenceAlignment.containsPossibleAlignment(englishPosition, frenchPosition)) {
			        		if (frenchPosition == 0) { // it's obsolete
			        			current = englishPosition;
			        			System.err.println("mother fucker");
			        			System.exit(0);
			        		} else {
			        			current = englishPosition;
			        			jumpCounter.incrementCount(current - previous, 1.0);
			        			previous = current;
			        		}
			        	}
			     }
//				 int positionIndex = positionIndexer.addAndGetIndex(new TIntArrayList(new int[]{englishLength, previous}));
			 }
		}
		
		
		private void modelOneTrainEStep(List<SentencePair> trainingSentencePairs) {
			for (SentencePair pair : trainingSentencePairs) {
				List<String> eList = pair.englishWords;
				List<String> fList = pair.frenchWords;
				innerTrain(eList, fList);
			}
		}
		
		private void innerTrain(List<String> eList, List<String> fList) {
			int numFrench = fList.size();
			int numEnglish = eList.size();
			
			for (int frenchPosition = 0; frenchPosition < numFrench; frenchPosition++) {
				int frenchIndex = frenchIndexer.indexOf(fList.get(frenchPosition));
				double tNull = tMap.getCount(NULL_INDEX, frenchIndex);
				double tSum = 0.0;
				for (int englishPosition = 0; englishPosition < numEnglish; englishPosition++) {
					int englishIndex = englishIndexer.indexOf(eList.get(englishPosition));
					tSum += tMap.getCount(englishIndex, frenchIndex);
				}
				double denorminator = (PNULL * tNull) + PREST / numEnglish * tSum;
				for (int englishPosition = 0; englishPosition < numEnglish; englishPosition++) {
					int englishIndex = englishIndexer.indexOf(eList.get(englishPosition));
					double pNormal = (PREST / numEnglish * tMap.getCount(englishIndex, frenchIndex))
							/ denorminator;
							//((PNULL * tNull) + PREST / numEnglish * tSum);
					counterMap.incrementCount(englishIndex, frenchIndex, pNormal);
				}
				double pSpecial = (PNULL * tNull)
								/ denorminator;
								//(PNULL * tNull + PREST / numEnglish * tSum);
				counterMap.incrementCount(NULL_INDEX, frenchIndex, pSpecial);
			}
		}
		
		public HMMFactory(List<SentencePair> validatePairs1,
				Map<Integer,Alignment> validateAlignments_1,
				List<SentencePair> validatePairs2,
				Map<Integer,Alignment> validateAlignments_2) {
			frenchCounter = new IntCounter();
			englishCounter = new IntCounter();
			counterMap = new IntCounterMap();
			tMap = new IntCounterMap();
			
			englishIndexer = new Indexer<String>();
			frenchIndexer = new Indexer<String>();
			
			jumpCounter = new IntCounter();
			positionIndexer = new Indexer<TIntList>();
			
			this.validatePairs1 = validatePairs1;
			this.validateAlignments_1 = validateAlignments_1;
			this.validatePairs2 = validatePairs2;
			this.validateAlignments_2 = validateAlignments_2;
		}
	}
	
	
	public Alignment alignSentencePair(SentencePair sentencePair) {
		
		Alignment alignment = new Alignment();
		int numFrenchWords = sentencePair.getFrenchWords().size();
		int numEnglishWords = sentencePair.getEnglishWords().size();
		int[][] viterbi = new int[numEnglishWords + 1][numFrenchWords];
		double[] highScores = new double[numEnglishWords + 1];
//		System.out.print("Pair # " + sentencePair.getSentenceID() + ": ");
		for (int frenchPosition = 0; frenchPosition < numFrenchWords; frenchPosition++) {
			int frenchIndex = frenchIndexer.indexOf(sentencePair.frenchWords.get(frenchPosition));
			for (int englishPosition = 0; englishPosition < numEnglishWords; englishPosition++) {
				if (frenchPosition == 0) {
//					viterbi[englishPosition][frenchPosition] = englishPosition;
					highScores[englishPosition] = 0;
					
				} else{
					resolveHighest(viterbi, highScores, frenchPosition, frenchIndex, numEnglishWords, sentencePair.englishWords);

				}
				
			}
//			alignment.addAlignment(maxPosition, frenchPosition, true);
		}
		//if (sentencePair.getSentenceID() < 8)
//			printViterbi(viterbi);
//		printHighScores(highScores);
		backwardResolve(alignment, viterbi, highScores);
		return alignment;
	}
	
	private void backwardResolve(Alignment alignment, int[][] viterbi, double[] highScores) {
		if (null == viterbi[0]) {
			System.err.println("shit");
			System.exit(-1);
		}
		int frenchLength = viterbi[0].length;
		int englishLength = viterbi.length;
		int[] reversedResult = new int[frenchLength]; 
			double max = Double.NEGATIVE_INFINITY;
			int maxPosition = -1;
			for (int englishPosition = 0; englishPosition < englishLength; englishPosition++) {
				if (highScores[englishPosition] > max) {
					max = highScores[englishPosition];
//					maxPosition = viterbi[englishPosition][frenchLength - 1];
					maxPosition = englishPosition;
				}
			}
			
			reversedResult[0] = maxPosition;
			maxPosition = viterbi[maxPosition - 1][frenchLength - 1];
			reversedResult[1] = maxPosition;
//			reversedResult[1] = viterbi[maxPosition - 1][frenchLength - 1];
			boolean output = false;
			
			for (int frenchPosition = 2; frenchPosition < frenchLength; frenchPosition++) {
				maxPosition = viterbi[maxPosition][frenchLength - frenchPosition];
				reversedResult[frenchPosition] = maxPosition;
			}
			
			for (int frenchPosition = frenchLength - 1; frenchPosition >= 0; frenchPosition--) { 
				alignment.addAlignment(reversedResult[frenchPosition], frenchLength - frenchPosition, true);
				if (output) {
					System.out.print(frenchLength - frenchPosition + ":" + reversedResult[frenchPosition] + "\t");
				}
			}
			if (output) {
				System.out.println();
			}
	}
	/**
	private void resolveHighest(int[][] viterbi, double[] highScores, int frenchPosition, int frenchIndex, int englishLength, List<String> englishWords) {
		
		for (int englishPosition = -1; englishPosition < englishLength; englishPosition++) {
			double max = Double.NEGATIVE_INFINITY;
			int maxPosition = -1;
			double denominator = 0.0;
			for (int i = -1; i < englishLength; i++) {
				denominator += jumpCounter.getCount(i - englishPosition);
			}
			for (int i = -1; i < englishLength; i++) {
				int englishIndex = -1;
				if (i >= 0) {
					englishIndex = englishIndexer.indexOf(englishWords.get(i));
				}
				double temp = Math.log(highScores[englishPosition+1])
						+ Math.log(jumpCounter.getCount(i - englishPosition) / denominator)
						+ Math.log(tMap.getCount(englishIndex, frenchIndex));
				///**
				if (Double.isNaN(Math.log(highScores[englishPosition+1]))) {
					System.err.println("first" + highScores[englishPosition+1]);
				}
				if (Double.isNaN(Math.log(jumpCounter.getCount(i - englishPosition) / denominator))) {
					System.err.println("second");
					System.out.println(denominator);
				}
				
				if (Double.isNaN(Math.log(tMap.getCount(englishIndex, frenchIndex)))) {
					System.err.println("third");

				}
				///
				
				if (temp > max) {
					max = temp;
					maxPosition = i;
				}
			}
			highScores[englishPosition + 1] = Math.exp(max);
			viterbi[englishPosition + 1][frenchPosition] = maxPosition;
//			System.out.println("maxPosition is " + maxPosition);
		}
	}
	*/
	
	private void resolveHighest(int[][] viterbi, double[] highScores, int frenchPosition, int frenchIndex, int englishLength, List<String> englishWords) {
		
		for (int englishPosition = 0; englishPosition < englishLength; englishPosition++) {
			double max = Double.NEGATIVE_INFINITY;
			int maxPosition = -1;
			double denominator = 0.0;
			for (int i = 0; i < englishLength; i++) {
				denominator += jumpCounter.getCount(i - englishPosition);
			}
			for (int i = 0; i < englishLength; i++) {
				int englishIndex = -1;
				if (i > 0) {
					englishIndex = englishIndexer.indexOf(englishWords.get(i));
				}
				double t1 = highScores[englishPosition]; 
					//Math.log(highScores[englishPosition]);
				double t2 = Math.log(jumpCounter.getCount(i - englishPosition) / denominator);
				double t3 = Math.log(tMap.getCount(englishIndex, frenchIndex));
//				double temp = Math.log(highScores[englishPosition])
//						+ Math.log(jumpCounter.getCount(i - englishPosition) / denominator)
//						+ Math.log(tMap.getCount(englishIndex, frenchIndex));
				double temp = t1
						+ t2
						+ t3;
				
				
				if (temp > max) {
					max = temp;
					maxPosition = i;
				}
				if (Double.isInfinite(temp) || Double.isNaN(temp)) {
					if (Double.isNaN(t1) || Double.isInfinite(t1)) {
//						System.err.println("first" + t1 + ", highScores is " + highScores[englishPosition] + ", englishPosition is " + englishPosition);
						
//						System.exit(0);
					}
					if (Double.isNaN(t2) || Double.isInfinite(t2)) {
//						System.err.println("second: " + t2 + ", " + (i - englishPosition) + ", "+ denominator);
					}
					
					if (Double.isNaN(t3) || Double.isInfinite(t3)) {
						System.err.println("third");

					}
				}
			}
			
//			highScores[englishPosition] = Math.exp(max);
			highScores[englishPosition] = max;
			if (Double.isInfinite(max) || Double.isNaN(max)) {
				System.out.println(highScores[englishPosition]);
				System.exit(0);
			}
			viterbi[englishPosition][frenchPosition] = maxPosition;
//			System.out.println("maxPosition is " + maxPosition);
		}
	}
	
	private void printViterbi(int[][] viterbi) {
		System.out.println("viterbi:");
//		for (int j = 0; j < viterbi[0].length; j++) {
//			System.out.print(j + "\t");
//		}
//		System.out.println();
		for (int i = 0; i < viterbi.length; i++) {
			if (viterbi[i] == null) {
				System.exit(-1);
			}
			
			
			for (int j = 0; j < viterbi[i].length; j++) {
				System.out.print(viterbi[i][j] + "\t");
			}
			System.out.println();
		}
//		System.exit(-1);
	}
	
	private void printHighScores(double[] highScores) {
		for (int i = 0; i < highScores.length; i++) {
			System.out.print(highScores[i] + "\t");
		}
		System.out.println();
//		System.exit(-1);
	}
	
	public HMMWordAligner(IntCounter frenchCounter, IntCounter englishCounter,
			IntCounterMap tMap, Indexer<String> englishIndexer,
			Indexer<String> frenchIndexer,
			IntCounter jumpCounter,
			Indexer<TIntList> positionIndexer) {
		this.frenchCounter = frenchCounter;
		this.englishCounter = englishCounter;
		this.tMap = tMap;
		this.englishIndexer = englishIndexer;
		this.frenchIndexer = frenchIndexer;
		
		this.jumpCounter = jumpCounter;
		this.positionIndexer = positionIndexer;

	}
	
	IntCounter frenchCounter;
	IntCounter englishCounter;
	IntCounterMap tMap;
	Indexer<String> englishIndexer;
	Indexer<String> frenchIndexer;
	
	IntCounter jumpCounter;
	Indexer<TIntList> positionIndexer;
}
